/*
 * Wazuh Vulnerability scanner - Scan Orchestrator
 * Copyright (C) 2015, Wazuh Inc.
 * March 4, 2024.
 *
 * This program is free software; you can redistribute it
 * and/or modify it under the terms of the GNU General Public
 * License (version 2) as published by the FSF - Free Software
 * Foundation.
 */

#ifndef _SCAN_INVENTORY_SYNC_HPP
#define _SCAN_INVENTORY_SYNC_HPP

#include "../policyManager/policyManager.hpp"
#include "chainOfResponsability.hpp"
#include "inventorySync.hpp"
#include "scanContext.hpp"

/**
 * @brief TScanInventorySync class.
 * This class is in charge of synchronizing the inventory database.
 * It receives the scan context and the inventory database and returns the scan context with the inventory updated.
 *
 * @tparam TScanContext scan context type.
 */
template<typename TScanContext = ScanContext>
class TScanInventorySync final
    : public AbstractHandler<std::shared_ptr<TScanContext>>
    , public InventorySync
{
public:
    // LCOV_EXCL_START
    /**
     * @brief ScanInventorySync constructor.
     *
     * @param inventoryDatabase Inventory database.
     */
    explicit TScanInventorySync(Utils::RocksDBWrapper& inventoryDatabase)
        : InventorySync(inventoryDatabase)
    {
    }
    // LCOV_EXCL_STOP

    /**
     * @brief Handles request and passes control to the next step of the chain.
     *
     * @param data Scan context.
     * @return std::shared_ptr<ScanContext> Abstract handler.
     */
    std::shared_ptr<TScanContext> handleRequest(std::shared_ptr<TScanContext> data) override
    {
        std::string key;
        key.append(data->agentNodeName());
        key.append("_");
        key.append(data->agentId());
        key.append("_");

        // Create the key for the inventory.
        key.append(affectedComponentKey(data));

        const auto& column = AFFECTED_COMPONENT_COLUMNS.at(data->affectedComponentType());
        std::vector<std::string> inventory;
        std::vector<std::string> inventoryBase;
        bool isInventoryEmpty = true;

        std::string value;
        if (m_inventoryDatabase.get(key, value, column))
        {
            inventory = Utils::split(value, ',');
            inventoryBase = inventory;
            isInventoryEmpty = false;
        }

        // Elements are the cves of the scan result.
        for (auto& [cve, value] : data->m_elements)
        {
            // If the cve is not in the inventory, insert it.
            if (std::find(inventory.begin(), inventory.end(), cve) == inventory.end())
            {
                std::string elementKey;
                elementKey.append(key);
                elementKey.append("_");
                elementKey.append(cve);

                data->m_elements[cve] = TInventorySync<TScanContext>::buildElement("INSERTED", elementKey);
                inventory.push_back(cve);
            }
        }

        // Remove elements from the inventory that are not in the scan result.
        inventory.erase(
            std::remove_if(inventory.begin(),
                           inventory.end(),
                           [&data, &key](const std::string& cve)
                           {
                               if (data->m_elements.find(cve) == data->m_elements.end())
                               {
                                   std::string elementKey;
                                   elementKey.append(key);
                                   elementKey.append("_");
                                   elementKey.append(cve);

                                   data->m_elements.emplace(
                                       cve, TInventorySync<TScanContext>::buildElement("DELETED", elementKey));
                                   logDebug2(
                                       WM_VULNSCAN_LOGTAG, "Removing element from inventory: %s", elementKey.c_str());
                                   return true; // Return true to remove this element from inventory
                               }
                               return false; // Otherwise, keep this element in inventory
                           }),
            inventory.end());

        // If the new scan have the element in the inventory, delete it from m_elements.
        auto it = data->m_elements.begin();
        while (it != data->m_elements.end())
        {
            // If the element is in the inventory after deleting the elements that are not in the scan result,
            // delete it from m_elements.
            if (std::find(inventoryBase.begin(), inventoryBase.end(), it->first) != inventoryBase.end())
            {
                if (std::find(inventory.begin(), inventory.end(), it->first) != inventory.end())
                {
                    it = data->m_elements.erase(it);
                }
                else
                {
                    ++it;
                }
            }
            else
            {
                ++it;
            }
        }

        // If the inventory is empty, set the initial scan to true.
        if (!m_inventoryDatabase.get(key, value, OS_INITIAL_SCAN))
        {
            // Put the initial value(utc time) for the initial scan.
            m_inventoryDatabase.put(key, Utils::getCurrentISO8601(), OS_INITIAL_SCAN);
        }
        else
        {
            data->m_isFirstScan = false;
        }

        if (inventory.empty())
        {
            // If the inventory previously had elements and now is empty, delete the inventory.
            if (!isInventoryEmpty)
            {
                logDebug2(WM_VULNSCAN_LOGTAG, "Deleting agent element key: %s", key.c_str());
                m_inventoryDatabase.delete_(key, column);
            }
        }
        else
        {
            std::sort(inventory.begin(), inventory.end());
            std::sort(inventoryBase.begin(), inventoryBase.end());

            // If the inventory has changed, update it.
            if (!std::equal(inventory.begin(), inventory.end(), inventoryBase.begin(), inventoryBase.end()))
            {
                std::string insertListString;
                for (const auto& cve : inventory)
                {
                    insertListString.append(cve);
                    insertListString.append(",");
                }
                insertListString.pop_back();
                logDebug2(WM_VULNSCAN_LOGTAG,
                          "Updating/Inserting agent element key: %s -> %s",
                          key.c_str(),
                          insertListString.c_str());
                m_inventoryDatabase.put(key, insertListString, column);
            }
            else
            {
                logDebug2(WM_VULNSCAN_LOGTAG, "No changes in agent element key: %s", key.c_str());
                return nullptr;
            }
        }

        return AbstractHandler<std::shared_ptr<TScanContext>>::handleRequest(std::move(data));
    }
};

using ScanInventorySync = TScanInventorySync<>;

#endif // _SCAN_INVENTORY_SYNC_HPP
